// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/dwio/unit_loader_test_tools.h>
#include <pollux/dwio/common/unit_loader_tools.h>

using kumo::pollux::dwio::common::LoadUnit;

namespace kumo::pollux::dwio::common::test {

ReaderMock::ReaderMock(
    std::vector<uint64_t> rowsPerUnit,
    std::vector<uint64_t> ioSizes,
    UnitLoaderFactory& factory,
    uint64_t rowsToSkip)
    : rowsPerUnit_{std::move(rowsPerUnit)},
      ioSizes_{std::move(ioSizes)},
      unitsLoaded_(std::vector<std::atomic_bool>(rowsPerUnit_.size())),
      loader_{factory.create(getUnits(), rowsToSkip)},
      currentUnit_{0},
      currentRowInUnit_{0} {
  POLLUX_CHECK(rowsPerUnit_.size() == ioSizes_.size());
  auto [currentUnit, currentRowInUnit] = unit_loader_tools::howMuchToSkip(
      rowsToSkip, rowsPerUnit_.cbegin(), rowsPerUnit_.cend());
  currentUnit_ = currentUnit;
  currentRowInUnit_ = currentRowInUnit;
}

bool ReaderMock::read(uint64_t maxRows) {
  if (!loadUnit()) {
    return false;
  }
  const auto rowsToRead =
      std::min(maxRows, rowsPerUnit_[currentUnit_] - currentRowInUnit_);
  loader_->onRead(currentUnit_, currentRowInUnit_, rowsToRead);
  currentRowInUnit_ += rowsToRead;
  return true;
}

void ReaderMock::seek(uint64_t rowNumber) {
  uint64_t totalRows = 0;
  uint64_t rowsLeft = rowNumber;
  for (size_t unit = 0; unit < rowsPerUnit_.size(); ++unit) {
    const uint64_t rowCount = rowsPerUnit_[unit];
    if (rowsLeft < rowCount) {
      currentUnit_ = unit;
      currentRowInUnit_ = rowsLeft;
      loader_->onSeek(currentUnit_, currentRowInUnit_);
      return;
    }
    rowsLeft -= rowCount;
    totalRows += rowCount;
  }
  POLLUX_CHECK_EQ(
      rowsLeft,
      0,
      "Can't seek to possition {} in file. Must be up to {}.",
      rowNumber,
      totalRows);
}

bool ReaderMock::loadUnit() {
  POLLUX_CHECK(currentRowInUnit_ <= rowsPerUnit_[currentUnit_]);
  if (currentRowInUnit_ == rowsPerUnit_[currentUnit_]) {
    currentRowInUnit_ = 0;
    ++currentUnit_;
    if (currentUnit_ >= rowsPerUnit_.size()) {
      return false;
    }
  }
  auto& unit = loader_->getLoadedUnit(currentUnit_);
  auto& unitMock = dynamic_cast<LoadUnitMock&>(unit);
  POLLUX_CHECK(unitMock.isLoaded());
  return true;
}

std::vector<std::unique_ptr<LoadUnit>> ReaderMock::getUnits() {
  std::vector<std::unique_ptr<LoadUnit>> units;
  for (size_t i = 0; i < rowsPerUnit_.size(); ++i) {
    units.emplace_back(std::make_unique<LoadUnitMock>(
        rowsPerUnit_[i], ioSizes_[i], unitsLoaded_, i));
  }
  return units;
}

std::vector<std::atomic_bool> getUnitsLoadedWithFalse(size_t count) {
  std::vector<std::atomic_bool> unitsLoaded(count);
  for (auto& unit : unitsLoaded) {
    unit = false;
  }
  return unitsLoaded;
}

} // namespace kumo::pollux::dwio::common::test
